from haystack.query import SearchQuerySet

from chatbot.bot.domain.question import Question
from chatbot.bot.domain.response import Response
from chatbot.bot.engine import ChatBot
from chatbot.bot.logic.logic_adapter import LogicAdapter
from chatbot.const import const
from chatbot.const.configuration import config
from chatbot.models import TableColumnValue, TableColumn
from chatbot.service.logic.ir.search_engine import get_result_from_search_engine
from chatbot.service.logic.table.detect_column_type import Detector
from chatbot.service.logic.table.table_data import TableData
from chatbot.settings import logger


class TableLogicAdapter(LogicAdapter):
    def __init__(self, chatbot: ChatBot, **kwargs):
        super().__init__(chatbot, **kwargs)
        self.response = Response()
        self.table_data = TableData().tables
        self.seg_words = []
        self.results = {}
        self.response_text = ''

    def can_process(self, question):
        query_text = question.text.strip()

        if not query_text:
            return False

        # Reject to process OCR
        if str(query_text).startswith('img'):
            return False

        self.seg_words = question.seg_text.split(const.SEG_SEPARATOR)

        all_corpus = get_result_from_search_engine(question.seg_text, (TableColumnValue, TableColumn))

        self.extract_column_and_value(all_corpus)

        # 获取比较信息，如A比B大多少
        compare = Detector().detect_operation(question)
        response_text = ''

        for table_id in self.results.keys():

            target_records = self.get_target_records(table_id)
            target_columns = self.extract_target_columns(table_id)

            self.generate_answer(table_id, target_columns, target_records, compare)

            if self.response_text:
                self.response = Statement(text=self.response_text)
                self.response.confidence = 1.0
            else:
                self.response = self.chatbot.get_unknown_answer(query_text)

            logger.info("query_text: {}, answer is: {} ".format(query_text, self.response.text))

            if self.response.confidence:
                return True

    def process(self, question: Question, additional_response_selection_parameters=None):
        return self.response

    def extract_column_and_value(self, qs: SearchQuerySet):
        """从搜索引擎得到的结果中提取表格的列和值

            提取时，需要大于搜索引擎返回结果的阈值 The :class:`Config`.SEARCH_ENGINE_CONFIDENCE

            :param qs: Haystack's SearchQuerySet
            :type: qs: Haystack's SearchQuerySet

            :returns: None
        """
        self.results = {}
        for search_result in qs:
            model_name = search_result.model_name
            score = search_result.score
            if score < config.SEARCH_ENGINE_CONFIDENCE:
                break

            obj = search_result.object
            table_id = obj.table.id

            if model_name == 'tablecolumn':
                column_id = obj.id
                column_name = obj.name
                value = None
                word = column_name

            # model_name = 'tablecolumnvalue'
            else:

                column_id = obj.col.id
                column_name = obj.col.name
                value = obj.value
                word = value

            # 控制列和值的精确匹配
            if not self.is_in_seg_words(word):
                continue

            self.set_results(table_id, column_id, column_name, value)

    def get_target_records(self, table_id: int) -> list:
        """基于提取的表格列和值推导表格查询的条件获取的结果数据

            :param table_id: table's id for search in pandas and self.results
            :type: int

            :returns: 基于pandas的Frames的target_records,
            :type: list
        """
        table = self.results.get(table_id)
        values = self.table_data[table_id]['values']
        target_records = []

        for column_id in table.keys():
            column = table.get(column_id)
            column_values = column.get('value', [])

            if column_values:
                for v in column_values:
                    target_value = values[values[column['column_name']] == v]
                    target_records.append(
                        {
                            'column': column['column_name'],
                            'value': v,
                            'target': target_value
                        }
                    )
        return target_records

    def extract_target_columns(self, table_id: int) -> list:
        """基于提取的表格列和值推导表格查询的目标列

        如果在表格列中没有提取到目标列，默认用表格全部列

            :param table_id: table's id for search in self.results
            :type: int

            :returns: 查询目标列的列名：target_columns,
            :type: list
        """
        table = self.results.get(table_id)
        target_columns = []

        for column_id in table.keys():
            column = table.get(column_id)
            column_name = column['column_name']
            value = column.get('value', [])

            if not value:
                target_columns.append(column_name)

        if not target_columns:
            target_columns = list(self.table_data[table_id]['columns'].keys())

        return target_columns

    def is_in_seg_words(self, word):
        # 控制列和值的精确匹配
        if word not in self.seg_words:
            return False
        return True

    def set_results(self, table_id, column_id, column_name, value):
        table = self.results.get(table_id, {})
        result = table.get(column_id, {})

        result['column_id'] = column_id
        result['column_name'] = column_name
        result['value'] = result.get('value', [])

        # 记录单独命中列的次数
        if not value:
            result['hit_num'] = result.get('hit_num', 0) + 1

        if value and value not in result['value']:
            result['value'].append(value)

        table[column_id] = result
        self.results[table_id] = table

    def generate_answer(self, table_id: int, target_columns: list, target_records: list, compare: list):
        """基于提取的目标列和查询条件，生成对应的答案

        :param table_id: table's id for search in self.results
        :type: int
        :param target_columns: 查询目标列
        :type: list
        :param target_records: 查询条件对应的结果数据
        :type: dict型的list
        :param compare: 比较相关的信息
        :type: dict型的list

        :returns: 无
        """
        self.response_text = ''
        has_compare = False if not len(compare) else True

        if not has_compare:
            self.generate_answer_with_no_compare(target_columns, target_records)
        else:
            self.generate_answer_with_compare(table_id, target_columns, target_records, compare)

    def generate_answer_with_no_compare(self, target_columns: list, target_records: list):
        """基于提取的目标列和查询条件，生成没有比较的query对应的答案

            :param target_columns: 查询目标列
            :type: list
            :param target_records: 查询条件对应的结果数据
            :type: dict型的list

            :returns: 无
        """
        for col in target_columns:
            for row in target_records:
                index_value = row['value']
                value = row['target']
                target_value = value[col].values
                self.response_text += ''.join(
                    [
                        "<p><span style='font-weight:bolder;'>{}</span> : {}</p>".format(index_value, v) for v in
                        target_value
                    ]
                )

    def generate_answer_with_compare(self, table_id: int, target_columns: list, target_records: list, compare: list):
        condition_len = len(target_records)
        target_len = len(target_columns)

        if condition_len == 2 and target_records[0]['column'] == target_records[1]['column']:
            # 处理A比B长多少的比较，同属性的比较
            self.process_a_compare_b_with_same_property(table_id, target_columns, target_records, compare)

        if condition_len == 1 and target_len == 1 or \
                condition_len == 2 and target_records[0]['column'] != target_records[1]['column']:
            # 处理比A长的有多少、比A长的B的C有多少
            self.process_a_compare_b_property(table_id, target_columns, target_records, compare)

    def process_a_compare_b_with_same_property(self, table_id: int, target_columns: list, target_records: list,
                                               compare: list):

        op = compare['op']
        slot_type = compare['slot_type']
        op_word = compare['word']
        unit = compare['unit']

        all_columns = self.table_data[table_id]['columns']

        if op in ('GT', 'LT'):
            # 处理A比B长多少的比较，同属性的比较
            for col in target_columns:
                column = all_columns[col]
                if column.col_type != slot_type:
                    continue

                greater = target_records[0]
                less = target_records[1]

                greater_index = greater['value']
                less_index = less['value']

                # 用正则方式提取值中的小数
                greater_value = self.extract_float(greater['target'][col].values[0])
                less_value = self.extract_float(less['target'][col].values[0])

                # 计算比较的两个值谁大谁小，保证greater大于less
                if greater_value < less_value:
                    greater_index, less_index = less_index, greater_index
                    greater_value, less_value = less_value, greater_value

                # 生成话术
                self.response_text += '{} 比 {} {}：{} {}'.format(greater_index, less_index, op_word,
                                                                round(greater_value - less_value, 2), unit)

    def process_a_compare_b_property(self, table_id: int, target_columns: list, target_records: list, compare: list):

        all_columns = self.table_data[table_id]['columns']
        column_values = self.table_data[table_id]['values']
        condition_len = len(target_records)

        # 处理比A长的有多少、比A长的B的C有多少
        target_col = target_columns[0]
        first_index_value = target_records[0]['value']

        condition_index_value, more_condition_col, more_condition_value, row = self.extract_more_condition(compare,
                                                                                                           condition_len,
                                                                                                           first_index_value,
                                                                                                           target_records)
        op = compare['op']
        operation_value = row[target_col].values[0]

        normalization = self.extract_float(operation_value)

        if op == 'LT':
            condition = column_values[target_col].map(
                lambda x: self.extract_float(x) < normalization)
        else:
            condition = column_values[target_col].map(
                lambda x: self.extract_float(x) > normalization)

        if condition_len == 1:
            target_values = column_values[condition]
        else:
            temp_values = column_values[condition]
            target_values = temp_values[temp_values[more_condition_col] == more_condition_value]

        self.response_text += "<p>{}的{}是{}</p>".format(
            condition_index_value, target_col, operation_value)

        self.response_text += "<p>总计有 : {} 条</p>".format(
            len(target_values))

        if not len(target_values):
            return

        self.response_text += '''<table class="layui-table">'''
        self.response_text += '<thead><tr>'

        final_col = []

        for col in all_columns.keys():
            if all_columns[col].is_primary_key:
                final_col.append(col)
                break
        final_col.append(target_col)

        for col in final_col:
            self.response_text += '<th>{}</th>'.format(col)

        self.response_text += '</tr></thead>'
        self.response_text += '<tbody>'

        for idx in target_values.index:
            self.response_text += '<tr>'
            for col in final_col:
                self.response_text += '<td>{}</td>'.format(target_values.loc[idx][col])
            self.response_text += '</tr>'
        self.response_text += '</tbody></table>'

    @staticmethod
    def extract_more_condition(condition_len, first_index_value, target_records):
        first_row = target_records[0]['target']
        first_col = target_records[0]['column']
        condition_index_value = first_index_value
        more_condition_col = first_col
        more_condition_value = first_index_value

        row = first_row
        if condition_len == 2:
            second_index_value = target_records[1]['value']
            second_row = target_records[1]['target']
            second_col = target_records[1]['column']

            if len(first_row) > len(second_row):
                row = second_row
                condition_index_value = second_index_value
                more_condition_col = first_col
                more_condition_value = first_index_value
            else:
                more_condition_col = second_col
                more_condition_value = second_index_value
        return condition_index_value, more_condition_col, more_condition_value, row

    @staticmethod
    def extract_float(text: str):
        import re
        return round(float(re.search(r'(\d+(\.\d+)?)', text).group()), 2)
